"""坐标系转换工具，目前仅包括等距圆柱投影的转换"""

import numpy as np
from math import pi, radians, cos, sin

radius = 1737.4 * 1000


class SphereCoordinate:
    """
    等距圆柱投影结果转球面坐标、平面坐标的工具
    """

    def __init__(
        self,
        modifier,
        lat1=60,
        lat2=-60,
        lon1=0,
        lon2=360,
        x_pixel_range=61440,
        y_pixel_range=184320,
    ):
        """
        Arguments:
            lat1 (float) : initial latitude in degree
            lat2 (float) : ending latitude in degree and should be less than lat1, e.g. lat1 = 60, lat2 = -60
            lon1 (float) : initial longitude in degree
            lon2 (float) : ending longitude in degree and should be larger than lon1, e.g. lon1 = 0, lon2 = 360
            x_pixel_range (int) : range of x pixel
            y_pixel_range (int) : range of y pixel
            modifier (float) : modifier
        """
        self.lat_range = lat1 - lat2
        self.lon_range = lon2 - lon1
        self.lat1 = lat1
        self.lon1 = lon1
        self.x_pixel_range = x_pixel_range
        self.y_pixel_range = y_pixel_range
        self.modifier = modifier

    def sphere2xyz(self, lat, lon, dr):
        """
        Warning ! The unit of latitude and longitude is degree.
        Arguments:
            lat (float) : latitude, from lat1 to lat2
            lon (float) : longitude, from lon1 to lon2
            dr (float) : radius in meter
        """
        lat = lat / 180 * pi
        lon = lon / 180 * pi
        x = (radius + dr) / self.modifier * np.cos(lat) * np.cos(lon)
        y = (radius + dr) / self.modifier * np.cos(lat) * np.sin(lon)
        z = (radius + dr) / self.modifier * np.sin(lat)
        return x, y, z

    def plane2sphere(self, x, y):
        """
        Return:
            lat (float) : latitude in degree, from lat1 to lat2
            lon (float) : longitude in degree, from lon1 to lon2
        """
        lat = self.lat1 - x * self.lat_range / self.x_pixel_range
        lon = y * self.lon_range / self.y_pixel_range + self.lon1
        return lat, lon

    def sphere2plane(self, lat, lon):
        """
        Return:
            x (float) : x in pixel
            y (float) : y in pixel
        """
        x = (self.lat1 - lat) * self.x_pixel_range / self.lat_range
        y = (lon - self.lon1) * self.y_pixel_range / self.lon_range
        return x, y

    def plane2xyz(self, x, y, dr):
        """
        Arguments:
            x (float) : x in pixel
            y (float) : y in pixel
            dr (float) : radius in meter
        """
        lat, lon = self.plane2sphere(x, y)
        return self.sphere2xyz(lat, lon, dr)

    def distortCircle(self, lat, d):
        """
        Arguments:
            lat (float) : latitude in degree, from lat1 to lat2
            d (float) : diameter in meter
        """
        D_x = d / radius * self.x_pixel_range / radians(self.lat_range)
        D_y = (
            d
            / radius
            * self.y_pixel_range
            / radians(self.lon_range)
            / cos(radians(lat))
        )
        return D_x, D_y

    def rotate(self, lat, lon):
        """
        三轴转角，将当前球面片旋转至天顶位置
        Arguments:
            lat (float) : latitude in degree
            lon (float) : longitude in degree
        """
        lat = radians(lat)
        lon = radians(lon)
        R_y = np.array(
            [
                [sin(lat), 0, -cos(lat)],
                [0, 1, 0],
                [cos(lat), 0, sin(lat)],
            ]
        )
        R_z = np.array(
            [
                [cos(lon), sin(lon), 0],
                [-sin(lon), cos(lon), 0],
                [0, 0, 1],
            ]
        )

        return R_y @ R_z

    def world2pixel(self, x):
        return x / self.modifier


class Projection:
    def __init__(self) -> None:
        pass

    def forward(self, lat, lon):
        raise NotImplementedError

    def inverse(self, x, y):
        raise NotImplementedError

    def update(self, **kwargs):
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)
        return self


class Equirectangular(Projection):
    def __init__(self, lat1, lon0, radius, pixel_size):
        """
        等距圆柱投影方式，虽然名义上投影按中心经线和标准纬线进行，实际上对应于图上时，中心经线等于起始经线，标准纬线等于起始纬线，平移量offset为0，这样才能保证投影前后，图像起始点与投影起始点一致\\
        Arguments:
            lat1 (float) : 标准纬线，单位为度
            lon0 (float) : 中心经线，单位为度
            radius (float) : 地球半径，单位为米
            pixel_size (float) : 像素大小，单位为米/像素
        """
        self.lat1 = lat1
        self.lon0 = lon0
        self.radius = radius
        self.pixel_size = pixel_size

    def forward(self, lat, lon):
        """
        Arguments:
            lat (float) : 纬度，单位为度
            lon (float) : 经度，单位为度
        """
        x = (self.lat1 - lat) / 180 * np.pi * self.radius / self.pixel_size
        y = (
            (lon - self.lon0)
            * np.cos(self.lat1 * np.pi / 180)
            * self.radius
            / 180
            / self.pixel_size
            * np.pi
        )
        return x, y

    def inverse(self, x, y):
        """
        Arguments:
            x (float) : 纵向坐标，单位为像素
            y (float) : 横向坐标，单位为像素
        """
        lat = self.lat1 - x * 180 / np.pi * self.pixel_size / self.radius
        lon = (
            y
            / np.cos(self.lat1 * np.pi / 180)
            * 180
            / np.pi
            * self.pixel_size
            / self.radius
            + self.lon0
        )
        return lat, lon


class SimpleCylinder(Projection):
    def __init__(self, lat0, lon0, lat_range, lon_range, x_size, y_size):
        """
        说明：
        Arguments:
            lat0 (float) : 中心纬度
            lon0 (float) : 中心经度
            lat_range (float) : 纬度范围
            lon_range (float) : 经度范围
            x_size (int) : 图像高度
            y_size (int) : 图像宽度
        """
        self.lat0 = lat0
        self.lon0 = lon0
        self.lat_range = lat_range
        self.lon_range = lon_range
        self.x_size = x_size
        self.y_size = y_size

    def forward(self, lat, lon):
        x = (self.lat0 - lat) * self.x_size / self.lat_range
        y = (lon - self.lon0) * self.y_size / self.lon_range
        return x, y

    def inverse(self, x, y):
        lat = self.lat0 - x * self.lat_range / self.x_size
        lon = y * self.lon_range / self.y_size + self.lon0
        return lat, lon


def rodrigues(rvec):
    rvec /= np.linalg.norm(rvec)
    location_vector = np.array([0, 0, 1])
    c = np.dot(rvec, location_vector)
    n_vector = np.cross(rvec, location_vector)

    n_vector_invert = np.array(
        (
            [0, -n_vector[2], n_vector[1]],
            [n_vector[2], 0, -n_vector[0]],
            [-n_vector[1], n_vector[0], 0],
        )
    )
    return (
        np.eye(3) + n_vector_invert + np.dot(n_vector_invert, n_vector_invert) / (1 + c)
    )
